package xjay.ai.emotionalsupport.rag;


import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Component
public class TarotRag extends RagService{
    private final String markdownResourcePath = "*.md";
    private final String pdfResourcePath = "*.pdf";
    public TarotRag(ResourcePatternResolver resolver, @Qualifier("pgVectorStore") VectorStore vectorStore) {
        String path = "classpath:document/tarot/";
        super.setMarkdownResourcePath(path + markdownResourcePath);
        super.setPdfResourcePath(path + pdfResourcePath);
        super.setResolver(resolver);
        super.setVectorStore(vectorStore);
    }

    public VectorStore rag() throws IOException {


        log.info("tarot rag：rag启动");
        try{
            if(getVectorStore().similaritySearch("*").isEmpty()) {
                getVectorStore().add(documents);
                log.info("tarot rag：rag启动完毕");
                return getVectorStore();
            }
        }   catch (Exception e) {
            log.error("向量数据库查询失败", e);
            throw e;
        }

        if (documents == null || documents.isEmpty()) {
            log.warn("documents列表为空，跳过映射创建");

        }
        // 预先建立文件名到文档的映射，提高性能
        Map<String, List<Document>> filename2Document = getDocuments().stream()
                .collect(Collectors.groupingBy(doc -> {
                    Map<String, Object> metadata = doc.getMetadata();
                    if (metadata == null) {
                        return "unknown";
                    }

                    // 先尝试获取filename，再尝试file_name
                    Object filename = metadata.get("filename");
                    if (filename == null) {
                        filename = metadata.get("file_name");
                    }

                    return filename != null ? filename.toString() : "unknown";
                }));

        List<Document> newDocuments = new ArrayList<>();
        log.info("tarot rag：resources size is {}", resources.size());

        for(Resource resource : resources) {
            String filename = resource.getFilename();

            if(filename != null) {

                //VectorStore的metadata查询语法: 通过SQL风格的条件表达式查询向量数据库中的文档。
                var existDocument = getVectorStore().similaritySearch("filename == '" + filename + "'");
                log.info("tarot rag: rag走到这里");
                 if(existDocument.isEmpty()) {

                     log.info("tarot rag：向量数据库添加新文档 {}", filename);

                     newDocuments.addAll(filename2Document.get(filename));
                 }
            }
        }
            // 添加到向量数据库
        if(!newDocuments.isEmpty()) {
            log.info("tarotRag 冷加载");
            getVectorStore().add(newDocuments);
        }
        log.info("tarot rag：rag启动完毕");
        return getVectorStore();
    }




}
